import torch
import torch.nn as nn
import torch.nn.functional as F


class DualChannelDetector(nn.Module):
    """
    双通道损伤检测网络：融合结构DNA与环境激励特征
    技术原理：CNN + 空间注意力机制
    刑侦比喻：资深警探交叉审问
    """

    def __init__(self, dna_dim=3, env_dim=1):
        """
        参数：
        dna_dim : DNA特征维度（频率/振型/阻尼）
        env_dim : 环境信号维度（时程数据）
        """
        super().__init__()

        # DNA分析通道（法医实验室）
        self.dna_conv = nn.Sequential(
            nn.Conv1d(dna_dim, 16, kernel_size=5),  # 滑动检测特征
            nn.ReLU(),
            nn.MaxPool1d(2)
        )

        # 环境分析通道（现场勘查）
        self.env_conv = nn.Sequential(
            nn.Conv1d(env_dim, 16, kernel_size=15),
            nn.ReLU(),
            nn.MaxPool1d(4)
        )

        # 注意力机制（聚焦关键证据）
        self.attention = nn.Sequential(
            nn.Linear(32, 16),
            nn.Tanh(),
            nn.Linear(16, 1),
            nn.Softmax(dim=1)
        )

        # 联合决策层（陪审团）
        self.jury = nn.Linear(32, 2)  # 输出损伤概率

    def forward(self, dna_data, env_data):
        """
        前向传播
        输入：
          dna_data : DNA特征 (batch, 3, seq_len)
          env_data : 环境信号 (batch, 1, seq_len)
        输出：
          verdict : 损伤概率 (batch, 2)
        """
        # 通道1：DNA特征提取
        dna_feat = self.dna_conv(dna_data)
        dna_feat = dna_feat.mean(dim=2)  # 全局平均池化

        # 通道2：环境特征提取
        env_feat = self.env_conv(env_data)
        env_feat = env_feat.mean(dim=2)

        # 特征拼接
        combined = torch.cat([dna_feat, env_feat], dim=1)

        # 注意力加权
        attn_weights = self.attention(combined)
        weighted_feat = attn_weights * combined

        # 最终判决
        verdict = self.jury(weighted_feat)
        return F.softmax(verdict, dim=1)